package org.example.service;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.client.WebSocketClient;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.util.concurrent.ListenableFuture;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

/**
 * WebSocket客户端，用于连接SenseVoice流式语音识别API
 */
public class SenseVoiceWebSocketClient {
    private static final Logger logger = LoggerFactory.getLogger(SenseVoiceWebSocketClient.class);

    private final String serverUrl;
    private final Consumer<String> onTextReceived;
    private final Consumer<String> onError;
    private WebSocketSession session;
    private boolean connected = false;

    // 连接状态标记
    private final AtomicBoolean connecting = new AtomicBoolean(false);

    // 未发送的消息队列，用于断线重连后恢复
    private final ConcurrentLinkedQueue<Object> pendingMessages = new ConcurrentLinkedQueue<>();

    // 会话状态
    private String sessionId;
    private boolean streaming = false;

    /**
     * 创建SenseVoice WebSocket客户端
     * 
     * @param serverUrl      SenseVoice WebSocket服务器URL
     * @param onTextReceived 接收文本消息的回调函数
     * @param onError        错误处理的回调函数
     */
    public SenseVoiceWebSocketClient(String serverUrl, Consumer<String> onTextReceived, Consumer<String> onError) {
        this.serverUrl = serverUrl;
        this.onTextReceived = onTextReceived;
        this.onError = onError;
        connect();
    }

    /**
     * 连接到SenseVoice WebSocket服务器
     */
    private void connect() {
        // 防止并发连接
        if (connecting.compareAndSet(false, true)) {
            try {
                logger.info("开始连接SenseVoice WebSocket服务: {}", serverUrl);
                WebSocketClient client = new StandardWebSocketClient();
                final CountDownLatch connectLatch = new CountDownLatch(1);

                WebSocketHandler handler = new TextWebSocketHandler() {
                    @Override
                    public void afterConnectionEstablished(WebSocketSession session) {
                        logger.info("SenseVoice WebSocket连接已建立：{}", serverUrl);
                        SenseVoiceWebSocketClient.this.session = session;
                        SenseVoiceWebSocketClient.this.connected = true;
                        connectLatch.countDown();

                        // 连接成功后处理待发送队列
                        processPendingMessages();
                    }

                    @Override
                    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
                        try {
                            String payload = message.getPayload();
                            logger.debug("接收到SenseVoice文本消息: {}", payload);
                            onTextReceived.accept(payload);
                        } catch (Exception e) {
                            logger.error("处理SenseVoice文本消息时出错", e);
                            onError.accept("处理文本消息时出错: " + e.getMessage());
                        }
                    }

                    @Override
                    public void handleTransportError(WebSocketSession session, Throwable exception) {
                        logger.error("SenseVoice WebSocket传输错误", exception);
                        onError.accept("WebSocket传输错误: " + exception.getMessage());
                        SenseVoiceWebSocketClient.this.connected = false;

                        // 尝试自动重连
                        scheduleReconnect();
                    }

                    @Override
                    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
                        logger.info("SenseVoice WebSocket连接已关闭: {}", status);
                        SenseVoiceWebSocketClient.this.connected = false;

                        // 如果不是正常关闭，尝试重连
                        if (status.getCode() != CloseStatus.NORMAL.getCode()) {
                            scheduleReconnect();
                        }
                    }
                };

                // 连接到WebSocket服务器
                ListenableFuture<WebSocketSession> future = client.doHandshake(handler, serverUrl);

                // 等待连接建立或超时
                boolean awaitComplete = connectLatch.await(5, TimeUnit.SECONDS);
                if (!awaitComplete) {
                    logger.error("连接SenseVoice WebSocket服务器超时: {}", serverUrl);
                    onError.accept("连接SenseVoice服务器超时");
                    this.connected = false;
                    future.cancel(true);

                    // 连接超时，尝试重连
                    scheduleReconnect();
                }
            } catch (Exception e) {
                logger.error("连接SenseVoice WebSocket服务器时出错: {}", e.getMessage());
                onError.accept("连接SenseVoice服务器时出错: " + e.getMessage());
                this.connected = false;

                // 连接异常，尝试重连
                scheduleReconnect();
            } finally {
                connecting.set(false);
            }
        } else {
            logger.info("已有连接正在进行中，忽略此次连接请求");
        }
    }

    /**
     * 安排一个延迟重连
     */
    private void scheduleReconnect() {
        if (!connecting.get()) {
            new Thread(() -> {
                try {
                    logger.info("计划在3秒后重新连接...");
                    Thread.sleep(3000); // 3秒后重连
                    connect();
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            }).start();
        }
    }

    /**
     * 处理待发送队列中的消息
     */
    private void processPendingMessages() {
        while (!pendingMessages.isEmpty() && isConnected()) {
            Object message = pendingMessages.poll();
            try {
                if (message instanceof TextMessageWrapper) {
                    TextMessageWrapper textMsg = (TextMessageWrapper) message;
                    session.sendMessage(new TextMessage(textMsg.getPayload()));
                    logger.debug("已从队列发送文本消息: {}", textMsg.getPayload());
                } else if (message instanceof BinaryMessageWrapper) {
                    BinaryMessageWrapper binaryMsg = (BinaryMessageWrapper) message;
                    session.sendMessage(new BinaryMessage(ByteBuffer.wrap(binaryMsg.getData())));
                    logger.debug("已从队列发送二进制消息，大小: {} 字节", binaryMsg.getData().length);
                }
            } catch (IOException e) {
                logger.error("从队列发送消息失败: {}", e.getMessage());
                // 发送失败，重新加入队列
                pendingMessages.add(message);
                break;
            }
        }
    }

    /**
     * 发送元数据到SenseVoice服务器
     * 
     * @param metadata 元数据JSON字符串
     */
    public void sendMetadata(String metadata) {
        if (!isConnected()) {
            // 如果未连接，加入待发送队列并尝试重连
            pendingMessages.add(new TextMessageWrapper(metadata));
            reconnect();
            return;
        }

        try {
            if (isConnected()) {
                session.sendMessage(new TextMessage(metadata));
                logger.debug("已发送元数据到SenseVoice: {}", metadata);

                // 更新流式状态
                if (metadata.contains("\"streaming\":true")) {
                    streaming = true;
                } else if (metadata.contains("\"streaming\":false")) {
                    streaming = false;
                }
            } else {
                logger.error("无法发送元数据，连接未建立");
                // 加入待发送队列并尝试重连
                pendingMessages.add(new TextMessageWrapper(metadata));
                reconnect();
            }
        } catch (IOException e) {
            logger.error("发送元数据到SenseVoice时出错: {}", e.getMessage());
            onError.accept("发送元数据时出错: " + e.getMessage());
            connected = false;

            // 加入待发送队列并尝试重连
            pendingMessages.add(new TextMessageWrapper(metadata));
            reconnect();
        }
    }

    /**
     * 发送音频数据到SenseVoice服务器
     * 
     * @param audioData 音频二进制数据
     * @param isFinal   是否为最终数据
     */
    public void sendAudioData(byte[] audioData, boolean isFinal) {
        if (!isConnected()) {
            // 如果未连接，加入待发送队列并尝试重连
            pendingMessages.add(new BinaryMessageWrapper(audioData, isFinal));
            reconnect();
            return;
        }

        try {
            if (isConnected()) {
                session.sendMessage(new BinaryMessage(ByteBuffer.wrap(audioData)));
                logger.debug("已发送音频数据到SenseVoice，大小: {} 字节, 是否最终: {}, 流式模式: {}",
                        audioData.length, isFinal, streaming);

                // 如果是最终数据且在非流式模式，重置状态
                if (isFinal && !streaming) {
                    // sessionId = null; // 保留sessionId以便后续使用
                }
            } else {
                logger.error("无法发送音频数据，连接未建立");
                // 加入待发送队列并尝试重连
                pendingMessages.add(new BinaryMessageWrapper(audioData, isFinal));
                reconnect();
            }
        } catch (IOException e) {
            logger.error("发送音频数据到SenseVoice时出错: {}", e.getMessage());
            onError.accept("发送音频数据时出错: " + e.getMessage());
            connected = false;

            // 加入待发送队列并尝试重连
            pendingMessages.add(new BinaryMessageWrapper(audioData, isFinal));
            reconnect();
        }
    }

    /**
     * 检查连接是否已建立
     * 
     * @return 连接状态
     */
    public boolean isConnected() {
        return connected && session != null && session.isOpen();
    }

    /**
     * 获取流式处理状态
     * 
     * @return 是否处于流式处理模式
     */
    public boolean isStreaming() {
        return streaming;
    }

    /**
     * 设置流式处理状态
     * 
     * @param streaming 是否启用流式处理
     */
    public void setStreaming(boolean streaming) {
        this.streaming = streaming;
    }

    /**
     * 重新连接到SenseVoice服务器
     */
    private void reconnect() {
        logger.info("正在重新连接到SenseVoice WebSocket服务器...");
        close();
        connect();
    }

    /**
     * 关闭连接
     */
    public void close() {
        if (session != null && session.isOpen()) {
            try {
                session.close();
            } catch (IOException e) {
                logger.error("关闭SenseVoice WebSocket连接时出错", e);
            }
        }
        connected = false;
        streaming = false;
    }

    /**
     * 文本消息包装类，用于队列存储
     */
    private static class TextMessageWrapper {
        private final String payload;

        public TextMessageWrapper(String payload) {
            this.payload = payload;
        }

        public String getPayload() {
            return payload;
        }
    }

    /**
     * 二进制消息包装类，用于队列存储
     */
    private static class BinaryMessageWrapper {
        private final byte[] data;
        private final boolean isFinal;

        public BinaryMessageWrapper(byte[] data, boolean isFinal) {
            this.data = data;
            this.isFinal = isFinal;
        }

        public byte[] getData() {
            return data;
        }

        public boolean isFinal() {
            return isFinal;
        }
    }
}